import os
import subprocess

class avclassWrapper():
    def __init__(self,laberScript,JsonPath,savePath):
        """ 
            执行该命令即可将vt报告转换成对应的家族标记
            简单解释:
                -c  能理解为只输出最可能的家族名,其余恶意软件类型  行为等tag不会输出
                -p  后面接输出的文件地址,将输出结果保存进去
                -vt 这个指样本是vt的json报告
        """
        if not os.path.exists(savePath):
            with open(savePath,"w") as file:
                file.close()
        self.cmd = "python3 {} -vt {} -vt3  -c -p > {}".format(laberScript, JsonPath, savePath)

    def runCmd(self):
        try:
            # 使用 stdout=subprocess.PIPE 和 stderr=subprocess.PIPE 来捕获输出和错误
            result = subprocess.run(self.cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
            
            # 打印标准输出
            print(result.stdout)
            
            # 如果有标准错误，打印标准错误
            if result.stderr:
                print("Error:", result.stderr)
            
            # 检查命令的返回码
            if result.returncode != 0:
                print(f"Command failed with return code {result.returncode}")

        except Exception as e:
            print(f"Exception occurred: {e}")

if __name__=="__main__":
    laberScript = r"/home/cyw/projects/malware_detected/find_sample_lable/avclass/avclass2/avclass2_labeler.py"
    jsonPath = r"/home/cyw/projects/malware_detected/saveData/origin_customer_malware_label/lableJson.txt"
    savePath = r"/home/cyw/projects/malware_detected/saveData/origin_customer_malware_label/sampleLables.txt"
    avclassTools = avclassWrapper(laberScript,jsonPath,savePath)
    avclassTools.runCmd()